use std::any::Any;
use std::collections::HashMap;
use std::sync::Arc;
use std::time::Duration;
use async_trait::async_trait;
use tokio::sync::Mutex;
use tokio::time::sleep;
use rand::Rng;

use protoactor_rs::actor::{Actor, ActorContext, ActorError, DefaultActorContext, PID, Props};
use protoactor_rs::system::ActorSystem;

// 基础的消息类型
#[derive(Debug, Clone)]
struct ChainMessage {
    message: String,
    hops_left: usize,
    path: Vec<String>,
}

#[derive(Debug, Clone)]
struct ChainResult {
    original_message: String,
    final_path: Vec<String>,
}

#[derive(Debug, Clone)]
struct Broadcast {
    message: String,
}

#[derive(Debug, Clone)]
struct BroadcastResponse {
    from: String,
    original_message: String,
}

// 简化的工作节点Actor
#[derive(Clone)]
struct WorkerActor {
    name: String,
    next_workers: Arc<Mutex<Vec<PID>>>,
}

#[async_trait]
impl Actor for WorkerActor {
    type Context = DefaultActorContext;

    async fn receive(&self, ctx: &mut Self::Context, msg: Box<dyn Any + Send>) -> Result<(), ActorError> {
        // 处理链式消息
        if let Some(chain_msg) = msg.downcast_ref::<ChainMessage>() {
            println!("[{}] Received chain message: {} (hops left: {})", 
                     self.name, chain_msg.message, chain_msg.hops_left);
            
            // 记录当前路径
            let mut new_path = chain_msg.path.clone();
            new_path.push(self.name.clone());
            
            if chain_msg.hops_left <= 1 {
                // 到达终点，发送结果
                println!("[{}] Chain complete, responding with result", self.name);
                let result = ChainResult {
                    original_message: chain_msg.message.clone(),
                    final_path: new_path,
                };
                ctx.respond(result);
            } else {
                // 继续链式传递
                let next_workers = self.next_workers.lock().await;
                
                if next_workers.is_empty() {
                    // 如果没有下一跳，直接返回结果
                    println!("[{}] No next hop available, ending chain early", self.name);
                    let result = ChainResult {
                        original_message: chain_msg.message.clone(),
                        final_path: new_path,
                    };
                    ctx.respond(result);
                } else {
                    // 随机选择一个下一跳
                    let idx = rand::thread_rng().gen_range(0..next_workers.len());
                    let next_pid = &next_workers[idx];
                    
                    println!("[{}] Forwarding chain message to next worker", self.name);
                    
                    // 构造新的消息
                    let new_msg = ChainMessage {
                        message: chain_msg.message.clone(),
                        hops_left: chain_msg.hops_left - 1,
                        path: new_path.clone(),
                    };
                    
                    // 尝试发送请求，如果失败则直接返回结果
                    match next_pid.request::<ChainResult>(new_msg, Duration::from_secs(5)).await {
                        Ok(result) => {
                            println!("[{}] Received result from next worker, forwarding", self.name);
                            ctx.respond(result);
                        },
                        Err(e) => {
                            println!("[{}] Error getting response from next worker: {:?}", self.name, e);
                            ctx.respond(ChainResult {
                                original_message: chain_msg.message.clone(),
                                final_path: new_path,
                            });
                        }
                    }
                }
            }
        }
        // 处理广播消息
        else if let Some(broadcast) = msg.downcast_ref::<Broadcast>() {
            println!("[{}] Received broadcast: {}", self.name, broadcast.message);
            
            // 构造响应
            let response = BroadcastResponse {
                from: self.name.clone(),
                original_message: broadcast.message.clone(),
            };
            
            ctx.respond(response);
        }
        else {
            println!("[{}] Received unknown message type", self.name);
        }
        
        Ok(())
    }
    
    async fn started(&self, _ctx: &mut Self::Context) -> Result<(), ActorError> {
        println!("[{}] Worker started", self.name);
        Ok(())
    }
    
    async fn stopping(&self, _ctx: &mut Self::Context) -> Result<(), ActorError> {
        println!("[{}] Worker stopping", self.name);
        Ok(())
    }
}

// 简化的管理器Actor
#[derive(Clone)]
struct ManagerActor {
    name: String,
    workers: HashMap<String, PID>,
}

#[async_trait]
impl Actor for ManagerActor {
    type Context = DefaultActorContext;

    async fn receive(&self, ctx: &mut Self::Context, msg: Box<dyn Any + Send>) -> Result<(), ActorError> {
        // 处理链式消息
        if let Some(chain_msg) = msg.downcast_ref::<ChainMessage>() {
            println!("[{}] Received chain command: {} (hops: {})", 
                     self.name, chain_msg.message, chain_msg.hops_left);
            
            if self.workers.is_empty() {
                println!("[{}] No workers available", self.name);
                ctx.respond(ChainResult {
                    original_message: chain_msg.message.clone(),
                    final_path: vec!["NO_WORKERS".to_string()],
                });
                return Ok(());
            }
            
            // 随机选择一个起始worker
            let workers_vec: Vec<&PID> = self.workers.values().collect();
            let idx = rand::thread_rng().gen_range(0..workers_vec.len());
            let start_pid = workers_vec[idx];
            
            println!("[{}] Starting chain with worker #{}", self.name, idx);
            
            // 发送请求并等待结果
            match start_pid.request::<ChainResult>(chain_msg.clone(), Duration::from_secs(10)).await {
                Ok(result) => {
                    println!("[{}] Chain complete, final path: {:?}", self.name, result.final_path);
                    ctx.respond(result);
                },
                Err(e) => {
                    println!("[{}] Error in chain: {:?}", self.name, e);
                    ctx.respond(ChainResult {
                        original_message: chain_msg.message.clone(),
                        final_path: vec!["ERROR".to_string()],
                    });
                }
            }
        }
        // 处理广播消息
        else if let Some(broadcast) = msg.downcast_ref::<Broadcast>() {
            println!("[{}] Broadcasting message: {}", self.name, broadcast.message);
            
            let mut responses = Vec::new();
            
            // 向所有worker广播消息
            for (worker_name, pid) in &self.workers {
                println!("[{}] Sending broadcast to {}", self.name, worker_name);
                
                match pid.request::<BroadcastResponse>(broadcast.clone(), Duration::from_secs(2)).await {
                    Ok(response) => {
                        println!("[{}] Received response from {}", self.name, response.from);
                        responses.push(response);
                    },
                    Err(e) => {
                        println!("[{}] Error getting response from {}: {:?}", self.name, worker_name, e);
                    }
                }
            }
            
            // 返回所有响应
            ctx.respond(responses);
        }
        else {
            println!("[{}] Received unknown message type", self.name);
        }
        
        Ok(())
    }
    
    async fn stopping(&self, _ctx: &mut Self::Context) -> Result<(), ActorError> {
        println!("[{}] Manager stopping", self.name);
        Ok(())
    }
}

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    println!("Starting actor communication example");
    
    // 初始化日志
    protoactor_rs::logging::init_global_logger(None, None)?;
    
    // 创建actor系统
    let system = ActorSystem::new().unwrap();
    println!("Actor system created");
    
    // 创建worker actors和它们的连接关系
    let worker_names = vec!["Worker1", "Worker2", "Worker3", "Worker4", "Worker5"];
    let connections = vec![
        ("Worker1", vec!["Worker2", "Worker3"]),
        ("Worker2", vec!["Worker3", "Worker4"]),
        ("Worker3", vec!["Worker5"]),
        ("Worker4", vec!["Worker5"]),
        ("Worker5", vec!["Worker1"]),
    ];
    
    // 创建所有worker，但还不设置连接
    let mut worker_actors = HashMap::new();
    let mut worker_pids = HashMap::new();
    
    println!("Creating workers...");
    
    // 第一步：创建所有worker实例和PID
    for name in &worker_names {
        let worker = WorkerActor {
            name: name.to_string(),
            next_workers: Arc::new(Mutex::new(Vec::new())),
        };
        
        worker_actors.insert(name.to_string(), worker.clone());
        
        let props = Props::from_producer(move || worker.clone());
        let pid = system.spawn::<WorkerActor>(props).await?;
        
        println!("Worker {} spawned with PID: {:?}", name, pid);
        worker_pids.insert(name.to_string(), pid);
        
        // 等待一小段时间，确保Actor已启动
        sleep(Duration::from_millis(10)).await;
    }
    
    println!("Setting up connections...");
    
    // 第二步：设置worker之间的连接
    for (from, to_list) in connections {
        if let Some(from_worker) = worker_actors.get(from) {
            let mut next_pids = Vec::new();
            
            for to in to_list {
                if let Some(to_pid) = worker_pids.get(to) {
                    next_pids.push(to_pid.clone());
                    println!("Added connection {} -> {}", from, to);
                }
            }
            
            // 直接设置连接，不使用消息
            let mut next_workers = from_worker.next_workers.lock().await;
            *next_workers = next_pids;
        }
    }
    
    // 创建管理器actor
    let manager = ManagerActor {
        name: "Manager".to_string(),
        workers: worker_pids.clone(),
    };
    
    let props = Props::from_producer(move || manager.clone());
    let manager_pid = system.spawn::<ManagerActor>(props).await?;
    println!("Manager spawned with PID: {:?}", manager_pid);
    
    // 等待Actor启动和连接设置完成
    sleep(Duration::from_millis(500)).await;
    
    // 测试链式消息
    println!("\n--- Testing Chain Messages ---");
    let chain_msg = ChainMessage {
        message: "Hello from chain".to_string(),
        hops_left: 3,
        path: vec![],
    };
    
    match system.request::<ChainResult>(&manager_pid, chain_msg, Duration::from_secs(10)).await {
        Ok(result) => {
            println!("Chain complete!");
            println!("Original message: {}", result.original_message);
            println!("Path taken: {:?}", result.final_path);
        },
        Err(e) => println!("Error in chain process: {:?}", e),
    }
    
    // 测试广播消息
    println!("\n--- Testing Broadcast Messages ---");
    let broadcast = Broadcast {
        message: "Hello everyone!".to_string(),
    };
    
    match system.request::<Vec<BroadcastResponse>>(&manager_pid, broadcast, Duration::from_secs(10)).await {
        Ok(responses) => {
            println!("Broadcast complete! Received {} responses:", responses.len());
            for response in responses {
                println!("  From {}: \"{}\"", response.from, response.original_message);
            }
        },
        Err(e) => println!("Error in broadcast process: {:?}", e),
    }
    
    // 确保在退出前等待所有消息处理完成
    sleep(Duration::from_secs(2)).await;
    println!("\nExample completed successfully");
    
    // 停止所有Actor，确保优雅退出
    for (name, pid) in worker_pids.iter() {
        println!("Stopping worker: {}", name);
        match system.stop(pid).await {
            Ok(_) => println!("Worker {} stopped successfully", name),
            Err(e) => println!("Error stopping worker {}: {:?}", name, e),
        }
    }
    
    // 停止管理器Actor
    match system.stop(&manager_pid).await {
        Ok(_) => println!("Manager stopped successfully"),
        Err(e) => println!("Error stopping manager: {:?}", e),
    }
    
    // 等待最终清理
    sleep(Duration::from_millis(500)).await;
    
    Ok(())
} 